package com.ajocer.springbootinit.rag;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import io.milvus.client.MilvusClient;
import io.milvus.client.MilvusServiceClient;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.grpc.DataType;
import io.milvus.param.*;
import io.milvus.param.collection.*;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.SearchResultsWrapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

/**
 * Milvus向量存储实现
 */
@Slf4j
@Component
public class MilvusVectorStore implements EmbeddingStore<TextSegment> {

    @Value("${milvus.collection}")
    private String collectionName;

    @Value("${milvus.dimension}")
    private int dimension;

    @Value("${milvus.index-type}")
    private String indexType;

    @Value("${milvus.metric-type}")
    private String metricType;

    @Autowired
    private MilvusServiceClient client;
    
    // 存储嵌入向量
    private final Map<String, Embedding> embeddings = new ConcurrentHashMap<>();
    
    // 存储文本段
    private final Map<String, TextSegment> segments = new ConcurrentHashMap<>();
    
    // ID生成器
    private final AtomicLong idGenerator = new AtomicLong();

    @PostConstruct
    public void init() {
        try {
            log.info("初始化Milvus向量存储");

            // 检查集合是否存在，不存在则创建
            HasCollectionParam hasCollectionParam = HasCollectionParam.newBuilder()
                    .withCollectionName(collectionName)
                    .build();
            R<Boolean> hasCollection = client.hasCollection(hasCollectionParam);
            if (hasCollection.getData()) {
                log.info("集合 {} 已存在", collectionName);
            } else {
                createCollection();
                log.info("集合 {} 创建成功", collectionName);
            }

            // 加载集合到内存
            LoadCollectionParam loadCollectionParam = LoadCollectionParam.newBuilder()
                    .withCollectionName(collectionName)
                    .build();
            client.loadCollection(loadCollectionParam);
            log.info("集合 {} 已加载到内存", collectionName);
        } catch (Exception e) {
            log.error("初始化Milvus向量存储失败", e);
        }
    }

    /**
     * 创建集合
     */
    private void createCollection() {
        // 创建字段
        FieldType idField = FieldType.newBuilder()
                .withName("id")
                .withDataType(DataType.Int64)
                .withPrimaryKey(true)
                .withAutoID(true)
                .build();

        FieldType embeddingField = FieldType.newBuilder()
                .withName("embedding")
                .withDataType(DataType.FloatVector)
                .withDimension(dimension)
                .build();

        FieldType textField = FieldType.newBuilder()
                .withName("text")
                .withDataType(DataType.VarChar)
                .withMaxLength(65535)
                .build();

        // 创建集合参数
        CreateCollectionParam createCollectionParam = CreateCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .withDescription("存储文本段和嵌入向量")
                .addFieldType(idField)
                .addFieldType(embeddingField)
                .addFieldType(textField)
                .build();

        // 创建集合
        client.createCollection(createCollectionParam);
    }

    @Override
    public String add(Embedding embedding) {
        String id = String.valueOf(idGenerator.incrementAndGet());
        embeddings.put(id, embedding);
        return id;
    }

    @Override
    public void add(String id, Embedding embedding) {
        embeddings.put(id, embedding);
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings) {
        List<String> ids = new ArrayList<>(embeddings.size());
        for (Embedding embedding : embeddings) {
            ids.add(add(embedding));
        }
        return ids;
    }

    @Override
    public String add(Embedding embedding, TextSegment textSegment) {
        String id = String.valueOf(idGenerator.incrementAndGet());
        embeddings.put(id, embedding);
        segments.put(id, textSegment);
        
        try {
            // 构建插入数据
            List<InsertParam.Field> fields = new ArrayList<>();
            fields.add(new InsertParam.Field("text", Collections.singletonList(textSegment.text())));
            fields.add(new InsertParam.Field("embedding", Collections.singletonList(embedding.vectorAsList())));
            
            InsertParam insertParam = InsertParam.newBuilder()
                    .withCollectionName(collectionName)
                    .withFields(fields)
                    .build();
            
            client.insert(insertParam);
        } catch (Exception e) {
            log.error("添加嵌入向量失败", e);
        }
        
        return id;
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings, List<TextSegment> textSegments) {
        if (embeddings.size() != textSegments.size()) {
            throw new IllegalArgumentException("嵌入向量和文本段数量不匹配");
        }
        
        List<String> ids = new ArrayList<>(embeddings.size());
        for (int i = 0; i < embeddings.size(); i++) {
            ids.add(add(embeddings.get(i), textSegments.get(i)));
        }
        return ids;
    }

    @Override
    public List<EmbeddingMatch<TextSegment>> findRelevant(Embedding referenceEmbedding, int maxResults, double minScore) {
        try {
            // 计算所有嵌入向量与参考向量的相似度
            List<EmbeddingMatch<TextSegment>> matches = new ArrayList<>();
            
            for (Map.Entry<String, Embedding> entry : embeddings.entrySet()) {
                String id = entry.getKey();
                Embedding embedding = entry.getValue();
                
                // 计算余弦相似度
                double similarity = cosineSimilarity(referenceEmbedding, embedding);
                
                // 过滤低于最小分数的结果
                if (similarity >= minScore) {
                    TextSegment segment = segments.get(id);
                    if (segment != null) {
                        // 使用工厂类创建EmbeddingMatch实例
                        EmbeddingMatch<TextSegment> match = EmbeddingMatchFactory.create(similarity, segment);
                        matches.add(match);
                    }
                }
            }
            
            // 按相似度降序排序
            matches.sort((a, b) -> Double.compare(b.score(), a.score()));
            
            // 限制结果数量
            return matches.stream()
                    .limit(maxResults)
                    .collect(Collectors.toList());
        } catch (Exception e) {
            log.error("查找相似内容失败", e);
            return Collections.emptyList();
        }
    }
    
    /**
     * 计算余弦相似度
     */
    private double cosineSimilarity(Embedding e1, Embedding e2) {
        List<Float> v1 = e1.vectorAsList();
        List<Float> v2 = e2.vectorAsList();
        
        if (v1.size() != v2.size()) {
            throw new IllegalArgumentException("向量维度不匹配");
        }
        
        double dotProduct = 0.0;
        double norm1 = 0.0;
        double norm2 = 0.0;
        
        for (int i = 0; i < v1.size(); i++) {
            dotProduct += v1.get(i) * v2.get(i);
            norm1 += Math.pow(v1.get(i), 2);
            norm2 += Math.pow(v2.get(i), 2);
        }
        
        // 避免除以零
        if (norm1 == 0 || norm2 == 0) {
            return 0.0;
        }
        
        return dotProduct / (Math.sqrt(norm1) * Math.sqrt(norm2));
    }

    /**
     * 关闭连接
     */
    @PreDestroy
    public void close() {
        // 不需要关闭client，因为它是由Spring管理的Bean
        log.info("Milvus向量存储已关闭");
    }
} 